package datastructure.threadpool;

import java.util.ArrayDeque;
import java.util.Queue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;

/**
 * 简单线程池实现
 */
public class SimpleThreadPool {
    // 任务队列
    private final Queue<Runnable> taskQueue;
    // 工作线程数组
    private final WorkerThread[] workers;
    // 核心线程数
    private final int corePoolSize;
    // 最大任务队列容量
    private final int maxQueueSize;
    // 线程池是否关闭
    private volatile boolean isShutdown;

    /**
     * 拒绝策略接口
     */
    @FunctionalInterface
    public interface RejectedExecutionHandler {
        void rejectedExecution(Runnable r, SimpleThreadPool executor);
    }

    // 默认拒绝策略：抛出异常
    private static final RejectedExecutionHandler defaultHandler =
            (r, executor) -> { throw new RejectedExecutionException("任务 " + r + " 被拒绝，线程池已关闭或队列已满"); };

    private final RejectedExecutionHandler handler;

    /**
     * 构造函数
     * @param corePoolSize 核心线程数
     * @param maxQueueSize 最大任务队列容量
     * @param handler 拒绝策略
     */
    public SimpleThreadPool(int corePoolSize, int maxQueueSize, RejectedExecutionHandler handler) {
        if (corePoolSize <= 0 || maxQueueSize <= 0) {
            throw new IllegalArgumentException("核心线程数和队列大小必须为正数");
        }
        this.corePoolSize = corePoolSize;
        this.maxQueueSize = maxQueueSize;
        this.taskQueue = new ArrayDeque<>(maxQueueSize);
        this.workers = new WorkerThread[corePoolSize];
        this.handler = handler != null ? handler : defaultHandler;
        this.isShutdown = false;

        // 初始化并启动工作线程
        for (int i = 0; i < corePoolSize; i++) {
            workers[i] = new WorkerThread();
            workers[i].start();
        }
    }

    /**
     * 提交任务到线程池
     * @param task 要执行的任务
     */
    public void execute(Runnable task) {
        if (isShutdown) {
            handler.rejectedExecution(task, this);
            return;
        }

        synchronized (taskQueue) {
            // 如果队列已满，执行拒绝策略
            if (taskQueue.size() >= maxQueueSize) {
                handler.rejectedExecution(task, this);
                return;
            }
            // 将任务加入队列
            taskQueue.add(task);
            // 唤醒一个等待的工作线程
            taskQueue.notify();
        }
    }

    /**
     * 关闭线程池
     */
    public void shutdown() {
        isShutdown = true;
        // 唤醒所有等待的工作线程，让它们检查关闭标志
        synchronized (taskQueue) {
            taskQueue.notifyAll();
        }
    }

    /**
     * 工作线程类，负责从任务队列中获取并执行任务
     */
    private class WorkerThread extends Thread {
        @Override
        public void run() {
            Runnable task;

            // 循环获取任务，直到线程池关闭且任务队列为空
            while (true) {
                synchronized (taskQueue) {
                    // 如果队列为空且线程池未关闭，则等待新任务
                    while (taskQueue.isEmpty() && !isShutdown) {
                        try {
                            taskQueue.wait();
                        } catch (InterruptedException e) {
                            Thread.currentThread().interrupt();
                            return;
                        }
                    }

                    // 如果线程池已关闭且任务队列为空，则退出
                    if (isShutdown && taskQueue.isEmpty()) {
                        break;
                    }

                    // 从队列中获取任务
                    task = taskQueue.poll();
                }

                try {
                    // 执行任务
                    if (task != null) {
                        task.run();
                    }
                } catch (Exception e) {
                    System.err.println("任务执行异常: " + e.getMessage());
                }
            }
        }
    }

    // 测试方法
    public static void main(String[] args) {
        // 创建线程池：3个核心线程，队列容量为10
        SimpleThreadPool pool = new SimpleThreadPool(3, 10, (r, executor) -> {
            System.out.println("任务 " + r + " 被拒绝执行，队列已满");
        });

        // 提交20个任务
        for (int i = 0; i < 20; i++) {
            final int taskId = i;
            pool.execute(() -> {
                try {
                    // 模拟任务执行时间
                    TimeUnit.MILLISECONDS.sleep(100);
                    System.out.println("任务 " + taskId + " 由线程 " + Thread.currentThread().getName() + " 执行完成");
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }
            });
        }

        // 关闭线程池
        pool.shutdown();
    }
}
